#!/bin/bash

set -Eeuo pipefail
# set -x

cd $(dirname $(readlink -f $0))

#HOSTNAME=$(hostname)
#DISPLAY_ID=${DISPLAY#:}
#FULL_DISPLAY=$(xauth list | grep "unix" | awk '{print $1}')
#COOKIE=$(xauth list | grep "unix" | awk '{print $3}')

ROS_DISTRO=""
SSH_PORT=10022
GRPC_PORT=50052

SPEC_IMAGE=""
SPEC_CONTAINER_NAME=""
DOCKER_EXTRA_ARGS=""
COMMAND=""
while [[ $# -gt 0 ]]; do
    case "$1" in
        --ros-dist)
            ROS_DISTRO="$2"
            shift 2
            ;;
        --image)
            SPEC_IMAGE="$2"
            shift 2
            ;;
        --name)
            SPEC_CONTAINER_NAME="$2"
            shift 2
            ;;
        --ssh-port)
            SSH_PORT="$2"
            shift 2
            ;;
        --grpc-port)
            GRPC_PORT="$2"
            shift 2
            ;;
        --help)
            echo "Usage: $0 --ros-dist <name> --image <ros_image> [--name str] [--ssh-port port] [--grpc-port port] [--help] [restart] [-- extra docker args]"
            exit 0
            ;;
        --)
            shift
            DOCKER_EXTRA_ARGS+=("$@")
            break
            ;;
        restart)
            COMMAND="restart"
            shift
            ;;
        *)
            echo "Error: unknown argument '$1'"
            exit 1
            ;;
    esac
done

if [ -z "$ROS_DISTRO" ]; then
	echo "Error: specify ros distribution with --ros-dist <name>"
	exit 0
fi
XAUTH=/tmp/.docker.ros.${ROS_DISTRO}.desktop.xauth
if [ -z "$SPEC_IMAGE" ]; then
	echo "Error: specify at least one image"
	exit 1
fi
IMAGE_NAME=${SPEC_IMAGE}
CONT_NAME=${SPEC_CONTAINER_NAME:-ros-${ROS_DISTRO}-desktop-container}

# remove empty elements in extra args
declare -a DOCKER_EXTRA_ARGS_CLEAN=()
for arg in "${DOCKER_EXTRA_ARGS[@]}"; do
    if [[ -n "$arg" ]]; then
        DOCKER_EXTRA_ARGS_CLEAN+=("$arg")
    fi
done
# DOCKER_EXTRA_ARGS=("${DOCKER_EXTRA_ARGS[@]/''/}")

echo -e "\n------ Configurations ------"
echo -e "Command: ${COMMAND:-Default Start}"
echo -e "ROS Distro: $ROS_DISTRO"
echo -e "SSH Port: $SSH_PORT"
echo -e "GRPC Port: $GRPC_PORT"
echo -e "Image Name: $IMAGE_NAME"
echo -e "Container Name: $CONT_NAME"
echo -e "Docker Extra Args: ${DOCKER_EXTRA_ARGS_CLEAN[@]}"
echo -e "----------------------------\n"

touch $XAUTH
xauth nlist $DISPLAY | sed -e 's/^..../ffff/' | xauth -f $XAUTH nmerge -

if [ "$COMMAND" = "restart" ]; then
    echo "Restarting container $CONT_NAME..."
    docker restart $CONT_NAME
    docker exec -it $CONT_NAME bash
else
    docker run -it \
        -e DISPLAY=$DISPLAY \
        -e QT_X11_NO_MITSHM=1 \
        -e XAUTHORITY=$XAUTH \
        -v /tmp/.X11-unix:/tmp/.X11-unix \
        -v $XAUTH:$XAUTH:ro \
        -p $SSH_PORT:22 \
        -p $GRPC_PORT:50051 \
        "${DOCKER_EXTRA_ARGS_CLEAN[@]}" \
        --device /dev/dri \
        --name $CONT_NAME \
        $IMAGE_NAME bash -c "echo -e \"export DISPLAY=${DISPLAY};\nexport QT_X11_NO_MITSHM=1;\nexport XAUTHORITY=${XAUTH};\nexport MUJOCO_VER=3.3.4;\nexport MUJOCO_DIR=/root/workspace/mujoco/\\\${MUJOCO_VER};\nexport PATH=\\\${PATH}:/root/workspace/blender-app:\\\${MUJOCO_DIR}/bin;\nexport RMW_IMPLEMENTATION=rmw_cyclonedds_cpp;\n\" >> /root/.bashrc && /usr/sbin/sshd && /usr/bin/bash --rcfile /root/.bashrc"

fi

rm -f $XAUTH

